# Let's continue to update it to fit the meet.
# Wei updated 2012/05/17.

import sys
import cPickle as p

#10 secs to load it.
def loadIDF(objectFileName):
    print "loadIDF*******************"
    IDFDict = {}
    f =file(objectFileName, 'r')
    tuplesCouter = 0
    for currentLine in f.readlines():
        currentLineElements = currentLine.strip().split(" ")
        if len(currentLineElements) != 2:
            print "error:1"
            sys.exit(1)
        else:
            tuplesCouter += 1
            if tuplesCouter % 1000000 == 0:
                print "-->",tuplesCouter
            key = currentLineElements[0]
            if key not in IDFDict:
                IDFDict[key] = int(currentLineElements[1])
            else:
                print "unexpected duplicate"
                print "error:2"
                print "key",key
                sys.exit(1)

    print "Load IDFDict Job Done"
    print "IDFDict length:",len(IDFDict)
    return IDFDict

#20mins to load the whole thing, that is good enough.
def loadBM25Scores(objectFileName):
    print "loadBM25Scores*******************"
    BM25Dict = {}
    f =file(objectFileName, 'r')
    
    tuplesCouter = 0
    
    
    for currentLine in f.readlines():
        currentLineElements = currentLine.strip().split(" ")
        if len(currentLineElements) != 3:
            print "error:1"
            sys.exit(1)
        else:
            tuplesCouter += 1
            if tuplesCouter % 1000000 == 0:
                print "-->",tuplesCouter
            key = currentLineElements[0] + "_" + currentLineElements[1]
            if key not in BM25Dict:
                BM25Dict[key] = float(currentLineElements[2])
            else:
                print "unexpected duplicate"
                print "error:2"
                print "key",key
                sys.exit(1)
    
    print "Load BM25 Dict Job Done"
    print "BM25 Dict length:",len(BM25Dict)
    return BM25Dict

def dealWithQueryLevelFeature(inputFileName1):
    print "dealWithQueryLevelFeature*******************"
    return_query_Dict = {}
    inputFileHandler1 = open(inputFileName1, "r")
    #keyWordsDict = {}
    
    for currentLine in inputFileHandler1.readlines():
        currentLineElements = currentLine.split("\t")
        queryNumber = currentLineElements[0].strip()
        adhocTopicNumber = currentLineElements[2].strip()
        relevenceScore = currentLineElements[4].strip()
        queryInfoString = queryNumber + " " + adhocTopicNumber + " " + relevenceScore + " "
        
        queryContent = currentLineElements[1].strip()
        warcTrecID = currentLineElements[3].strip()
        
        returnTurple = (queryContent, warcTrecID)
        if returnTurple not in return_query_Dict:
            return_query_Dict[returnTurple] = queryInfoString
        else:
            print "Error1"
            print returnTurple
            sys.exit(1)
            
    
    return return_query_Dict
    #The key is the thing we hunt for.
      
    '''   
        for currentTerm in queryContentElement:
            key = currentTerm + "_" + warcTrecID
            if key not in keyWordsDict:
                keyWordsDict[key] = 1
            else:
                keyWordsDict[key] += 1
    
    keyWordsDict
    sorted_keyWordsDict = sorted(keyWordsDict.iteritems(), key=operator.itemgetter(0))
    
    print len(keyWordsDict)
    print type(sorted_keyWordsDict)
    
    
    for currentElement in sorted_keyWordsDict:
        (queryTerm_warcID,freq) = currentElement
        print queryTerm_warcID
        #print "freq:",freq
    '''
    inputFileHandler1.close()

def dealWithDocumentLevelFeature( dict0, queryDict,inputFileName2 ):
    print "dealWithDocumentLevelFeature*******************"
    #print len(queryDict)
    #print queryDict
    
    documentDict = {}
    NOTRecordedCounter = 0
    inputFileHandler2 = open(inputFileName2, "r")
    for currentLine in inputFileHandler2.readlines():
        currentLineElements = currentLine.strip().split(" ")
        if len(currentLineElements) == 9:
            
            warcTrecID = currentLineElements[0]
            if warcTrecID in dict0:
                url = currentLineElements[1]
                hamProbability = currentLineElements[3]
                isPopular = currentLineElements[4]    
                urlLength = currentLineElements[5]
                documentSizeInBytes = currentLineElements[6]
                numberOfWords = currentLineElements[7]
                numberOfDistantWords = currentLineElements[8]
                documentInfoString = url + " " + hamProbability + " " + isPopular + " " + urlLength + " " + documentSizeInBytes + " " +  numberOfWords + " " + numberOfDistantWords + " "; 
                
                docID = currentLineElements[2]
                if (warcTrecID,docID) not in documentDict:
                    documentDict[ (warcTrecID,docID) ] = documentInfoString
                else:
                    print "Error1"
                    sys.exit(1)
            else:
                pass
        else:
            NOTRecordedCounter += 1
                
    inputFileHandler2.close()
    print "# of documents NOT recorded:",NOTRecordedCounter
    print "# of documents been put into the documentDict:", len(documentDict)
    '''
    for eachElement in documentDict:
        print eachElement,documentDict[eachElement]
    '''
    '''
    for eachElement in queryDict:
        print eachElement,queryDict[eachElement]    
    '''
    print "queryDict:",len(queryDict)
    print "documentDict:",len(documentDict)
    
    return_query_document_Dict_counter = 0
    #print documentDict
    
    #Let's construct a query-document dict
    return_query_document_Dict = {}
    
    #This 2 for loop, kind of expensive.
    for currentElement in queryDict:
        (queryContent,warcTrecID) = currentElement
        
        for eachWARCIDDocIDElement in documentDict:
            (currentWARCID,currentDocID) = eachWARCIDDocIDElement
            
            if warcTrecID == currentWARCID:
                newTurple = (queryContent,currentDocID)
                
                if newTurple not in return_query_document_Dict:
                    return_query_document_Dict[newTurple] = queryDict[currentElement] + documentDict[ (warcTrecID,currentDocID) ]
                    return_query_document_Dict_counter += 1
                    print return_query_document_Dict_counter
                else:
                    print "Error2"
                    print newTurple
                    sys.exit(1)

    #print len(return_query_document_Dict)
    #print return_query_document_Dict
    return return_query_document_Dict

def dealWithTermLevelFeature( query_document_Dict , inputFileName3):
    print "dealWithTermLevelFeature*******************"
    #print len(query_document_Dict)
    #print query_document_Dict
    
    
    
    termDict = {}
    
    
    inputFileHandler3 = open(inputFileName3, "r")
    for currentLine in inputFileHandler3.readlines():
        currentLineElements = currentLine.strip().split(" ")
        word = currentLineElements[0]
        docID = currentLineElements[1]
        position = currentLineElements[2]
        context = currentLineElements[3]
        BM25Score = currentLineElements[4]
        key = (word, docID, BM25Score)
        if key not in termDict:
            termDict[key] = [(position,context)]
        else:
            #print termDict[key]
            termDict[key] = termDict[key] + [(position,context)]
            
    inputFileHandler3.close()
    #print len(termDict)
    
    
    #for key in termDict:
    #    print key,termDict[key]
    #print termDict[("day","55")]
    #print termDict[("day","7")]
    
    
    #Let's construct a query_document_term dict
    return_query_document_term_Dict = {}
    
    noPositionCounter = 0
    for currentElement in query_document_Dict:
        #Copy all the info from the previous dict first.
        if currentElement not in return_query_document_term_Dict:
            return_query_document_term_Dict[currentElement] = query_document_Dict[currentElement]
        else:
            print "Error3"
            print currentElement
            sys.exit(1)
        
        #copy all the term level from the term dict to the return dict.
        (queryContent,docID) = currentElement
        queryContentElements = queryContent.split(" ")
        for queryTerm in queryContentElements:
            eachTermAndIDFlag = False
            for eachTuple in termDict:
                (word,wordDocID,wordBM25) = eachTuple
                if word == queryTerm and wordDocID == docID:
                    
                    eachTermAndIDFlag = True
                    
                    termInfoString = str(termDict[eachTuple]) + " " + wordBM25 + " "
                    
                    return_query_document_term_Dict[currentElement] += termInfoString
                    #query_document_Dict[currentElement] += termInfoString
                    '''
                    if (queryTerm, docID) in termDict:
                        newTurpleString = "(" + queryTerm + "," + docID + "," + str( len(termDict[(queryTerm, docID)]) ) + ","
                        for eachPositionInfo in termDict[(queryTerm, docID)][: len(termDict[(queryTerm, docID)])-1]:
                            newTurpleString += eachPositionInfo + ","
                        newTurpleString += termDict[(queryTerm, docID)][-1]
                        newTurpleString += ")"
                        #print newTurpleString
                        #newTurple = tuple(newTurpleString)
                        
                        if newTurpleString not in return_query_document_term_Dict:
                            return_query_document_term_Dict[newTurpleString] = query_document_Dict[currentElement]
                        
                    else:
                        #print "NO Key:",queryTerm, docID
                        noPositionCounter += 1
                        newTurpleString = "(" + queryTerm + "," + docID + "," + str( 0 ) + ")"
                        #newTurple = tuple(newTurpleString)
                        if newTurpleString not in return_query_document_term_Dict:
                            return_query_document_term_Dict[newTurpleString] = query_document_Dict[currentElement]
                        else:
                            print "Error3"
                            print newTurpleString
                            sys.exit(1)
                    '''
            if not eachTermAndIDFlag:
                #fake position context list and fake BM25 score.
                termInfoString = str([]) + " " + str(0) + " "
                return_query_document_term_Dict[currentElement] += termInfoString
                
                noPositionCounter += 1
                print "No Position and Context (queryTerm, docID) pair in termDict:",queryTerm,docID
                
                
    
    print "Word appeared in the query but No Position and Context info in document Counter:",noPositionCounter
    
    return return_query_document_term_Dict
    
def loadNeededDocuments(inputFileName1):
    inputFileForUseEvaluatedDocumentsInfoFileHandler = open(inputFileName1, "r")
    finalDocumentDict = {}
        
    for currentLine in inputFileForUseEvaluatedDocumentsInfoFileHandler.readlines():
        currentLineElements = currentLine.split("\t")
        WARCTRECID = currentLineElements[3]
        if WARCTRECID not in finalDocumentDict:
            finalDocumentDict[WARCTRECID] = 1
        else:
            finalDocumentDict[WARCTRECID] += 1
    inputFileForUseEvaluatedDocumentsInfoFileHandler.close()
        
    
    return finalDocumentDict

def constructWekaTrainingFile(dict2,BM25Dict,IDFDict):
    print "constructWekaTrainingFile*******************"
    outputFeatureFileName = "/data1/team/weijiang/machine-learning-project-related/featuresGenerated/finalTrainingFiles/feature_file_machine_generated.arff"
    
    outputQueryFeatureFileHandler = open(outputFeatureFileName,"w")
    
    outputQueryFeatureFileHandler.write("@relation ranking" + "\n")
    outputQueryFeatureFileHandler.write("@attribute doc_id numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute query_id numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute query_terms numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute query_content string" + "\n")
    outputQueryFeatureFileHandler.write("@attribute doc_words numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute doc_distinct_words numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute doc_size numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute doc_waterloo_spam numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute url_length numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute url_in_top_million {1, 0}" + "\n")
    
    outputQueryFeatureFileHandler.write("@attribute term1BM25ScorePerDoc numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute term2BM25ScorePerDoc numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute term3BM25ScorePerDoc numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute term4BM25ScorePerDoc numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute term5BM25ScorePerDoc numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute term6BM25ScorePerDoc numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute termsBM25ScoreSumPerDoc numeric" + "\n")
    
    outputQueryFeatureFileHandler.write("@attribute term1IDF numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute term2IDF numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute term3IDF numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute term4IDF numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute term5IDF numeric" + "\n")
    outputQueryFeatureFileHandler.write("@attribute term6IDF numeric" + "\n")
    
    outputQueryFeatureFileHandler.write("@attribute relevence_score_label {-2, 0, 1, 2, 3}" + "\n")
    outputQueryFeatureFileHandler.write("@data" + "\n")

    NotInKeyCounter = 0
    
    for queryWARCTRECIDDocIDTupleElement in dict2:
        # init value for the BM25 score part
        
        termsBM25ScoreSumPerDoc = 0.0
        
        term1BM25ScorePerDoc = str(0)
        term2BM25ScorePerDoc = str(0)
        term3BM25ScorePerDoc = str(0)
        term4BM25ScorePerDoc = str(0)
        term5BM25ScorePerDoc = str(0)
        term6BM25ScorePerDoc = str(0)
        
        # init value for the IDF part.
        term1IDF = str(0)
        term2IDF = str(0)
        term3IDF = str(0)
        term4IDF = str(0)
        term5IDF = str(0)
        term6IDF = str(0)
        
        currentOutputLine = ""
        (queryContent,docID) = queryWARCTRECIDDocIDTupleElement
        
        #Do process with the queryContent.
        for index, currentTerm in enumerate( queryContent.strip().split(" ") ):
            #print index,currentTerm
            if index == 0:
                # for the IDF part.
                if currentTerm in IDFDict:
                    term1IDF = str(IDFDict[currentTerm])
                else:
                    print currentTerm,"Not in dict"
                    print "Error4"
                    sys.exit(1)
                    
                # for the BM25Score part.
                lookUpKey = currentTerm + "_" + docID
                
                if lookUpKey in BM25Dict:
                    term1BM25ScorePerDoc = str(BM25Dict[lookUpKey])
                    termsBM25ScoreSumPerDoc += BM25Dict[lookUpKey]
                else:
                    #print lookUpKey,"Not in dict"
                    NotInKeyCounter += 1
                    #print "Error4"
                    #sys.exit(1)
                                      
            if index == 1:
                if currentTerm in IDFDict:
                    term2IDF = str(IDFDict[currentTerm])
                else:
                    print currentTerm,"Not in dict"
                    print "Error4"
                    sys.exit(1)
                    
                # for the BM25Score part.
                lookUpKey = currentTerm + "_" + docID
                
                if lookUpKey in BM25Dict:
                    term2BM25ScorePerDoc = str(BM25Dict[lookUpKey])
                    termsBM25ScoreSumPerDoc += BM25Dict[lookUpKey]
                else:
                    #print lookUpKey,"Not in dict"
                    NotInKeyCounter += 1
                    #print "Error4"
                    #sys.exit(1)
                                
            if index == 2:
                if currentTerm in IDFDict:
                    term3IDF = str(IDFDict[currentTerm])
                else:
                    print currentTerm,"Not in dict"
                    print "Error4"
                    sys.exit(1)
                
                # for the BM25Score part.
                lookUpKey = currentTerm + "_" + docID
                
                if lookUpKey in BM25Dict:
                    term3BM25ScorePerDoc = str(BM25Dict[lookUpKey])
                    termsBM25ScoreSumPerDoc += BM25Dict[lookUpKey]
                else:
                    #print lookUpKey,"Not in dict"
                    NotInKeyCounter += 1
                    #print "Error4"
                    #sys.exit(1)
            if index == 3:
                if currentTerm in IDFDict:
                    term4IDF = str(IDFDict[currentTerm])
                else:
                    print currentTerm,"Not in dict"
                    print "Error4"
                    sys.exit(1)

                # for the BM25Score part.
                lookUpKey = currentTerm + "_" + docID
                
                if lookUpKey in BM25Dict:
                    term4BM25ScorePerDoc = str(BM25Dict[lookUpKey])
                    termsBM25ScoreSumPerDoc += BM25Dict[lookUpKey]
                else:
                    #print lookUpKey,"Not in dict"
                    NotInKeyCounter += 1
                    #print "Error4"
                    #sys.exit(1)
            if index == 4:
                if currentTerm in IDFDict:
                    term5IDF = str(IDFDict[currentTerm])
                else:
                    print currentTerm,"Not in dict"
                    print "Error4"
                    sys.exit(1)
                
                # for the BM25Score part.
                lookUpKey = currentTerm + "_" + docID
                
                if lookUpKey in BM25Dict:
                    term5BM25ScorePerDoc = str(BM25Dict[lookUpKey])
                    termsBM25ScoreSumPerDoc += BM25Dict[lookUpKey]
                else:
                    #print lookUpKey,"Not in dict"
                    NotInKeyCounter += 1
                    #print "Error4"
                    #sys.exit(1)
            if index == 5:
                if currentTerm in IDFDict:
                    term6IDF = str(IDFDict[currentTerm])
                else:
                    print currentTerm,"Not in dict"
                    print "Error4"
                    sys.exit(1)
                
                # for the BM25Score part.
                lookUpKey = currentTerm + "_" + docID
                
                if lookUpKey in BM25Dict:
                    term6BM25ScorePerDoc = str(BM25Dict[lookUpKey])
                    termsBM25ScoreSumPerDoc += BM25Dict[lookUpKey]
                else:
                    #print lookUpKey,"Not in dict"
                    NotInKeyCounter += 1
                    #print "Error4"
                    #sys.exit(1)

        queryContent = "'" + queryContent + "'"
        elements = dict2[queryWARCTRECIDDocIDTupleElement].strip().split(" ")
        queryID = elements[0]
        queryLength = str( len( queryContent.strip().split(" ") ) )
        #adhoc, all is 0
        topicNumber = elements[1]
        relevenceScore = elements[2]
        urlContent = elements[3]
        hamProbability = elements[4]
        isPopular = elements[5]
        URLLength = elements[6]
        docSizeInBytes = elements[7]
        docWords = elements[8]
        docDistantWords = elements[9]
        

        
        currentOutputLine = docID + "," + queryID + "," + queryLength + "," + queryContent + "," + docWords +  "," + docDistantWords + ","+ docSizeInBytes + ","+ hamProbability + "," + URLLength+ "," + isPopular + "," + term1BM25ScorePerDoc + "," + term2BM25ScorePerDoc + "," + term3BM25ScorePerDoc + "," + term4BM25ScorePerDoc + ","+ term5BM25ScorePerDoc + "," + term6BM25ScorePerDoc + ","+ str( termsBM25ScoreSumPerDoc ) + "," + term1IDF + ","+ term2IDF + "," + term3IDF + ","+ term4IDF + ","+ term5IDF + ","+ term6IDF + ","+ relevenceScore
        outputQueryFeatureFileHandler.write(currentOutputLine + "\n")    
        #print currentOutputLine
    outputQueryFeatureFileHandler.close()
    print "NotInKeyCounter:",NotInKeyCounter

#main program
inputFileName1 = "/data1/team/weijiang/machine-learning-project-related/featuresGenerated/query_features_2011_adhoc.txt"

dict0 = loadNeededDocuments(inputFileName1)
print "# of documents taken into evaluation:",len(dict0)

'''
for currentElement in dict0:
    print currentElement,dict0[currentElement]
'''

dict1 = dealWithQueryLevelFeature(inputFileName1)
'''
print len(dict1)
for queryWARCTRECIDTupleElement in dict1:
    print queryWARCTRECIDTupleElement,dict1[queryWARCTRECIDTupleElement]
'''

objectFileName = "interDict.txt"
'''
inputFileName2 = "/data1/team/weijiang/machine-learning-project-related/featuresGenerated/document_features.txt"
dict2 = dealWithDocumentLevelFeature( dict0, dict1,inputFileName2 )
print len(dict2)

f =file(objectFileName, 'w')
p.dump(dict2, f)
'''

f = file(objectFileName, 'r')  
storedDict = p.load(f)  
#print storedDict


#print dict2["('indiana child support', '6752')"]


'''
testDictForDict2 = {}
key1Part1 = "illinois state tax"
key1Part2 = "402"
key2Part1 = "interview thank you"
key2Part2 = "17076"
value1 = "142 0 -2 http://adfxadvisory.com/jmfyu/zaptg/refund/refund.htm 0 0 53 70445 9182 1858"
value2 = "119 0 0 http://www.celebnextdoor.com/celebs/Jamelia/ 4 0 44 68679 3818 1183"
testDictForDict2[ (key1Part1,key1Part2) ] = value1
testDictForDict2[ (key2Part1,key2Part2) ] = value2
'''

BM25Dict = {}
IDFDict = {}

BM25FileName = "/data1/team/weijiang/machine-learning-project-related/featuresGenerated/term_features_BM25_partial.txt"
BM25Dict = loadBM25Scores(BM25FileName)

#print BM25Dict

IDFFileName = "/data1/team/weijiang/machine-learning-project-related/featuresGenerated/term_features_IDF_partial.txt"
IDFDict = loadIDF(IDFFileName)

constructWekaTrainingFile(storedDict,BM25Dict,IDFDict)



'''
for queryWARCTRECIDDocIDTupleElement in dict2:
    print queryWARCTRECIDDocIDTupleElement,dict2[queryWARCTRECIDDocIDTupleElement]
'''

'''
inputFileName3 = "/data1/team/weijiang/machine-learning-project-related/featuresGenerated/term_features.txt"
dict3 = dealWithTermLevelFeature( storedDict, inputFileName3 )
print len(dict3)
for queryDocIDTermTupleElement in dict3:
    print queryDocIDTermTupleElement,dict3[queryDocIDTermTupleElement]
'''

#The code should be cleaned.
    
    
    
